using DataFrames
using QuadGK
using Distributions
using CSV
using LinearAlgebra
using JLD
using DelimitedFiles
using SparseArrays

clearconsole()

#-------------------------------------------------------------
# IRF Configuration
#-------------------------------------------------------------

H = 60 # maximum horizon
n_every = 50 # use every n_every'th draw from the posterior
sh_size = 3   # shock size, in multiples of standard deviations
sh_id = 1 # sh_id = 1: TFP shock

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
include(readDir *"vech.jl");
include(readDir *"logSpline_Procedures.jl");
include(readDir *"VAR_MoreLags_Procedures.jl");
include(readDir *"Loaddata.jl");

#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nDataSel    = "2"
nVARSpec    = "4P" # with 1= percentiles, and  2= gini, frac
nVARMCMCSpec= "1"
specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/AltVARspec" * nVARSpec * ".jl")
include(specDir * "/AltVARMCMCspec" * nVARMCMCSpec * ".jl")

#-------------------------------------------------------------
# load aggregate data
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/Para" * nDataSel *"/"
agg_data = loadaggdata(1,Tend)
n_agg    = size(agg_data)[2]
percentiles_data, vec_percs, gini_coef_data, probMass_below1_data, probMass_above4_data = loadaltdata(1,Tend)

#-------------------------------------------------------------
# Combine agg data and density coefficients
#-------------------------------------------------------------
if nSpecStr == "Pctl"
    data    = [ agg_data percentiles_data ]
    n_cross = size(percentiles_data)[2]
else
    data    = [ agg_data gini_coef_data probMass_above4_data]
    n_cross = 2
end

n = n_agg + n_cross
p = nlags

#-------------------------------------------------------------
# Load posterior draws
#-------------------------------------------------------------
sName    = "AltVAR" * nVARSpec * "_MCMC" * nVARMCMCSpec * "_" * nSpecStr
loadDir  = "$(pwd())/results/Para" *nDataSel * "/" * sName *"/";

Bpdraw     = load(loadDir * sName * "_PostDraws.jld", "Bpdraw")
Sigtrpdraw = load(loadDir * sName * "_PostDraws.jld", "Sigtrpdraw")
Bpmean     = load(loadDir * sName * "_PostMeans.jld", "Bpmean")
Sigtrpmean = load(loadDir * sName * "_PostMeans.jld", "Sigtrpmean")

#-------------------------------------------------------------
# Generate IRFs at posterior mean of (Phi,Sigmatr)
#-------------------------------------------------------------

println("")
println("Generating IRFs at posterior mean... ")
println("")

YY_IRF         = zeros(H+1,n_agg+n_cross)

Sigtilde = Symmetric(Sigtrpmean*Sigtrpmean')
L = cholesky(Sigtilde).L
Btilde = reshape(Bpmean, n*p,n)
response = IRFredu(Btilde,L,H,sh_size,sh_id)
YY_IRF[2:end,:] = response'

savedir = "$(pwd())/results/Para" *nDataSel * "/" * sName *"/";
try mkdir(savedir) catch; end
CSV.write(savedir * sName * "_IRF_YY_AggSh"*string(sh_id)*"_pmean.csv", DataFrame(YY_IRF))

#-------------------------------------------------------------
# Generate IRFs for a subset of posterior draws
#-------------------------------------------------------------

println("")
println("Generating IRFs for posterior draws... ")
println("")

nsim = size(Bpdraw)[1]
n_subseq                 = floor(Int,nsim/n_every)
YY_IRF_uncertainty       = zeros(H+1,n,n_subseq)

# hh=1 is steady state
# hh=2 is period of impact

for pp = 1:n_subseq

    time_init_loop = time_ns();
    println(" Draw number:  $pp")
    println(" Remaining draws:  $(n_subseq-pp)")
    println(" Posterior draw number: $(pp*n_every)")

    Sigtilde = Symmetric(Sigtrpdraw[pp*n_every,:,:]*Sigtrpdraw[pp*n_every,:,:]')
    L = cholesky(Sigtilde).L
    Btilde = Bpdraw[pp*n_every,:,:]
    response = IRFredu(Btilde,L,H,sh_size,sh_id)
    YY_IRF_uncertainty[2:end,:,pp] = response'

    CSV.write(savedir * sName * "_IRF_YY_AggSh" * string(sh_id) * "_" * string(pp) * ".csv", DataFrame(YY_IRF_uncertainty[:,:,pp]))

    time_loop=signed(time_ns()-time_init_loop)/1000000000
    println(" Elapsed time = $(time_loop)")
    println("")

end
